from keras.layers import Lambda
from keras.layers import Concatenate
from keras.layers import Input, Dense
from keras.models import Model
from keras.layers import SeparableConv2D
from keras.layers import Conv1D
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import AveragePooling2D, GlobalAveragePooling2D
from keras.layers import Activation
from keras.layers import BatchNormalization
from keras.layers import DepthwiseConv2D
from keras.layers import concatenate
from keras.layers import Reshape
from keras.layers.core import Flatten
from keras import backend as K



class SelfModel(object):

    def build(self, input_shape = (256, 256, 3)):
        inputs = Input(shape=input_shape)
        x = SeparableConv2D(16, (3, 3), dilation_rate=1, padding = 'same')(inputs) # 主要利用空洞卷积
        x = MaxPooling2D(pool_size=2, padding='same')(x)

        # x = self.Inception(x, 32)
        x = self.Inception(x, 30)
        x = MaxPooling2D(pool_size=(3, 3), strides=(2, 2), padding='same')(x)

        # x = Conv2D(64, 1, padding = 'same')(x)
        # x = MaxPooling2D(pool_size = 3, strides = 2, padding = 'same')(x)
        # x = Conv2D(128, 3, padding='same')(x)
        # x = MaxPooling2D(pool_size=3, strides=2, padding='same')(x)

        x_1, x_2 = self.channel_split(x)

        x_1 = Conv2D(filters=64, kernel_size=1, strides=1, padding='same',
                   use_bias=False)(x_1)
        x_1 = BatchNormalization(axis=-1)(x_1)
        x_1 = Activation('relu')(x_1)

        x_1 = DepthwiseConv2D(kernel_size=3, strides=1, padding='same',
                            use_bias=False)(x_1)
        x_1 = BatchNormalization(axis=-1)(x_1)
        x_1 = Conv2D(filters=64, kernel_size=1, strides=1, padding='same',
                   use_bias=False)(x_1)
        x_1 = BatchNormalization(axis=-1)(x_1)
        x_1 = Activation('relu')(x_1)

        x = Concatenate(axis=-1)([x_1, x_2])
        x = Lambda(self.channel_shuffle)(x)

        x = AveragePooling2D(pool_size=3, strides=2, padding='same')(x)

        x = Flatten()(x) # 这里向量为?,? ，不知是否有问题？？？？？？

        # x = Dense(1024, activation='relu')(x)
        # x = Dense(512, activation='relu')(x)
        x = Dense(128, activation='relu')(x)
        y = Dense(4, activation='softmax')(x)

        model = Model(inputs = inputs, outputs = y)
        return model

    def channel_split(self, x):
        in_channels = x.shape.as_list()[-1] # 将最后一维数值拿出，拆分成两份
        ip = in_channels // 2
        # lambda是一种函数的写法，https://www.cnblogs.com/caizhao/p/7905094.html
        # c_hat 与 c 各取一半通道的特征图（最后一维）
        c_hat = Lambda(lambda z: z[:, :, :, 0:ip])(x)
        c = Lambda(lambda z: z[:, :, :, ip:])(x)

        return c_hat, c


    def channel_shuffle(self, x):
        height, width, channels = x.shape.as_list()[1:]  # 取出图片尺寸和通道数
        channels_per_split = channels // 2  # 拆分两份
        # -1代表不确定，如果整体点数确定。reshape函数 应该 会自动根据后面的参数，最后分配-1处的具体值
        # 下面三步将通道以2为一组，均匀分配（因为这里拆分成两份，决定了一组内通道间的间隔数）
        # 如[0 1 2 3 4 5 6 7]  转为 [0 4 1 5 2 6 3 7]
        x = K.reshape(x, [-1, height, width, 2, channels_per_split])
        x = K.permute_dimensions(x, (0,1,2,4,3))
        x = K.reshape(x, [-1, height, width, channels])

        return x

    def Conv2d_BN(self, x, nb_filter, kernel_size, padding='same', strides=(1, 1), name=None):
        if name is not None:
            bn_name = name + '_bn'
            conv_name = name + '_conv'
        else:
            bn_name = None
            conv_name = None

        x = Conv2D(nb_filter, kernel_size, padding=padding, strides=strides, activation='relu', name=conv_name)(x)
        x = BatchNormalization(axis=3, name=bn_name)(x)
        return x


    def Inception(self, x, nb_filter):
        branch1x1 = self.Conv2d_BN(x, nb_filter, (1, 1), padding='same', strides=(1, 1), name=None)

        branch3x3 = self.Conv2d_BN(x, nb_filter, (1, 1), padding='same', strides=(1, 1), name=None)
        branch3x3 = self.Conv2d_BN(branch3x3, nb_filter, (3, 3), padding='same', strides=(1, 1), name=None)

        branch5x5 = self.Conv2d_BN(x, nb_filter, (1, 1), padding='same', strides=(1, 1), name=None)
        branch5x5 = self.Conv2d_BN(branch5x5, nb_filter, (1, 1), padding='same', strides=(1, 1), name=None)

        branchpool = MaxPooling2D(pool_size=(3, 3), strides=(1, 1), padding='same')(x)
        branchpool = self.Conv2d_BN(branchpool, nb_filter, (1, 1), padding='same', strides=(1, 1), name=None)

        x = concatenate([branch1x1, branch3x3, branch5x5, branchpool], axis=3)

        return x


if __name__ == '__main__':
    model = Model()
    model.build()